/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import alloy.ast.BinaryExpr;
import alloy.ast.BinaryExprOp;
import alloy.ast.Command;
import alloy.ast.CompOp;
import alloy.ast.ComprehensionExpr;
import alloy.ast.Decl;
import alloy.ast.DeclIter;
import alloy.ast.Decls;
import alloy.ast.ElemFormula;
import alloy.ast.EmptyFormula;
import alloy.ast.Expr;
import alloy.ast.Formula;
import alloy.ast.Id;
import alloy.ast.Module;
import alloy.ast.ModuleInfo;
import alloy.ast.Multiplicity;
import alloy.ast.Node;
import alloy.ast.QuantifiedFormula;
import alloy.ast.Quantifier;
import alloy.ast.SetMultExpr;
import alloy.ast.SigExpr;
import alloy.ast.Signature;
import alloy.ast.TypedExpr;
import alloy.ast.VarCreator;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.type.RelationType;
import alloy.type.UnionType;

/* NOTE:  this will not work properly before atomization.  Even after atomization, the new
semantics might not be compatible with the analysis performed here */

/**
 * QuantifierRemover transforms an AST such that all quantified formulas of the
 * form "all x : A | x.r = f(x)" become "r = {a : A, b : B, c : C ... | b->c->... in f(a))" where
 * A->B->C->... is the type of r.
 **/
public class QuantifierRemover extends EqDepthFirstReplacer {
    private ModuleInfo moduleInfo;

    /** Mapping from Variable Ids to SigExpr representing the variable types **/
    private Map removableVars;

    /** Quantifier-free equality formulas extracted from the AST **/
    private Set extractedFormulas;
    
    private QuantifierRemover(ModuleInfo moduleInfo_) {
        moduleInfo = moduleInfo_;
	removableVars = new HashMap();
	extractedFormulas = new HashSet();
    }
    
    public static Node removeQuantifiers(Node node, Command command) {
	//System.out.println("\n*Removing quantifiers for " + command);
	QuantifierRemover qrem = new QuantifierRemover(((Module)command.getParent().getParent()).getModuleInfo());
	node = (Node) node.applyReturnVisitor(qrem);
	for (Iterator iter = qrem.extractedFormulas.iterator(); iter.hasNext();) {
	    node = ((Formula) node).and((Formula) iter.next());
	}

	return node;
    }
  
 
    /**
     * @param node node whose leaves are to be examined
     * @param varSet the set of quantified variables <code>node</code> is allowed to contain
     * @return true if any leaves reachable from the input node 
     * are unbound quantified variables; otherwise returns false
     **/
    private boolean _hasUnboundVars(Node node, Set varSet) {
	
	if (node instanceof VariableExpr) {
	    //System.out.println("  analyzing VarExpr " + node);
	    //System.out.println("  in varSet? " +
	    //		       varSet.contains(((VariableExpr) node).getVar().getId()));
	    //System.out.print("  field? ");
	    try {
		moduleInfo.getSigsForField(node.nodeString());
		//System.out.println("true");
		return false;
	    } 
	    catch (ModuleInfo.NoSigsException nse) {
		//System.out.println("false");
		return !varSet.contains(((VariableExpr) node).getVar().getId());
	    }
	    
		    
	}
	
	
	if (node instanceof VarCreator) {
	    for (DeclIter varIter = new DeclIter((VarCreator) node); varIter.hasNext();) { 
		varIter.next();
		varSet.add(varIter.getVariable().getId());
	    }
	    return _hasUnboundVars(((VarCreator) node).getBody(), varSet);
	    
	}
	
	boolean hasUVars = false;
	for (int i = 0; i < node.numChildren() && !hasUVars; i++) 
	    hasUVars = _hasUnboundVars((Node) node.childAt(i), varSet);
	
	return hasUVars;
    }

    /**
     * @param binExpr binary expression whose form is to be examined
     * @return true if <code>expr</code> is an expression of the form
     * x.r where x is a universally quantified variable in this.removableVars 
     * and r is a field of SigExpr representing the type of x.  Otherwise returns false.
     **/
    private boolean _isTransformableDotExpr(BinaryExpr expr) {
		if (expr.getOp().equals(BinaryExprOp.JOIN_DOT) ||
		        expr.getOp().equals(BinaryExprOp.JOIN_DOUBLECOLON) ||
		        expr.getOp().equals(BinaryExprOp.JOIN_SQBRACKETS))  {
		    try {
			Id vId = ((VariableExpr) expr.getLeft()).getVar().getId();
			if (removableVars.containsKey(vId)) {
			    String vType = ((SigExpr) removableVars.get(vId)).nodeString();
			    for (Iterator sigIter = moduleInfo.getSigsForField(expr.getRight().nodeString()).iterator();
				 sigIter.hasNext(); ) { 
				if (((Signature) sigIter.next()).toSigExpr().nodeString().equals(vType))
				  return true;
			    }
			}
		    }
		    catch (ClassCastException cce) {}
		    catch (ModuleInfo.NoSigsException nse) {}
		}
	
		return false;
    }

    /**
     * If the universally quantified variable can be factored out of input expression,
     * the transformed expression is added to this.extractedFormulas and true is returned.
     * @param expr1 must be a transformable dot expression if true is to be returned
     * @param expr2 cannot contain any externally bound quantified variables except the one
     * in <code>expr1</code> if true is to be returned
     * @return true if a transformation is possible, false otherwise.
     **/
    private boolean _transform(Expr expr1, Expr expr2) {

	BinaryExpr dotExpr = ((expr1 instanceof BinaryExpr) ? (BinaryExpr) expr1 : null);

	if ((dotExpr != null) && _isTransformableDotExpr(dotExpr)) {
	    
	    // one side of the quantified formula is of the form x.r where r is a relation;
	    // now, we constrict the other side to have no other quantified, externally bound 
	    // variables besides x
	    Set usableVars = new HashSet();
	    usableVars.add(((VariableExpr) dotExpr.getLeft()).getVar().getId());
	    //System.out.println("  testing if " + expr2 + " has unbound vars.");
	   	    
	    if (!_hasUnboundVars(expr2, usableVars)) {
		
		//System.out.println("  success: " + usableVars);
				
		RelationType rType = ((TypedExpr) dotExpr.getRight()).getType().toRelationType();
		/** the type of the dot expression is a union type which cannot be reduced to 
		 a relation type**/
		if (rType == null) return false;
		
		// construct one variable for each basic type in rType:
		// a : A, b : B, c : C . . . where rType = A->B->C->...
		Decls decls = new Decls();
		VariableExpr[] vExprs = new VariableExpr[rType.numBasicTypes()];
		
		for (int i = 0; i < rType.numBasicTypes(); i++) {
		    Expr typeExpr = rType.getBasicTypeAt(i).toExpr();
		    Variable var = new Variable(Id.generateId(typeExpr.nodeString()));
		    vExprs[i] = new VariableExpr(var);
		    vExprs[i].setType(new UnionType(new RelationType(rType.getBasicTypeAt(i))));
		    SetMultExpr sme = new SetMultExpr(Multiplicity.NONE, (Expr) typeExpr.copy());
		    sme.setType(new UnionType(new RelationType(rType.getBasicTypeAt(i))));
		    decls.addDecl(Decl.in((Variable) var.copy(),sme));
		}
		
		// construct the "b->c->..." part of the "in" formula
		Expr rangeExpr = (Expr) vExprs[1].copy();
		for (int i = 2; i < vExprs.length; i ++)
		    rangeExpr = rangeExpr.product((Expr) vExprs[i].copy());
		
		// construct the "f(a)" part of the "in" formula
		Map varExprMapping = new HashMap();
		varExprMapping.put(dotExpr.getLeft().nodeString(), vExprs[0].copy());
		expr2 = (Expr) expr2.
		    applyReturnVisitor(new ExprAndTypeReplaceVisitor(varExprMapping,"", ""));
		//System.out.println("   a = " + vExprs[0]);
		//System.out.println("   f(a) = " + expr2);

		// construct the comprehension
		ComprehensionExpr compExpr = 
		    new ComprehensionExpr(decls, 
					  new ElemFormula(rangeExpr, CompOp.SUBSETIN, expr2), 
					  true);
		compExpr.setType(new UnionType(rType));
		
		// construct "r = {a:A, b:B, c:C, ... | b->c->... in f(a)}"
		extractedFormulas.add(new ElemFormula((Expr) dotExpr.getRight().copy(),
						      CompOp.EQUALS,
						      compExpr));

		//System.out.println("  Replacing " + expr1 + " = " + expr2 + " with " +
		//		   new ElemFormula((Expr) dotExpr.getRight().copy(),
		//				   CompOp.EQUALS,
		//				   compExpr));  
		return true;
	    }
	    
	    
	}
	
	return false;
    }

    public Object visit(QuantifiedFormula quantifiedFormula_) {
		// formula must be universally quantified and its body must be
		// an elementary "=" formula
	
		if (quantifiedFormula_.getQuantifier().equals(Quantifier.ALL)) {
		    // extract variables which are universally quantified over an entire signature,
		    // and store them in this.removableVars; otherwise, store extracted vars in
		    // this.unremovableVars
		    for (DeclIter varIter = new DeclIter(quantifiedFormula_); varIter.hasNext();) { 
				varIter.next();
				if (varIter.getMultiplicity() instanceof SetMultExpr) {
				    SetMultExpr var = (SetMultExpr) varIter.getMultiplicity();
				    if (((var.getMult().equals(Multiplicity.ONE)) ||
					 (var.getMult().equals(Multiplicity.NONE))) &&
					(var.getExpr() instanceof SigExpr)) {
					//System.out.println("  Adding variable " + varIter.getVariable().getId() + " of type " +
					//		   var.getExpr());
					removableVars.put(varIter.getVariable().getId(), var.toExpr());
				    }
				}
		    }
		}
		return super.visit((Node) quantifiedFormula_);
    }

    public Object visit(ElemFormula elemFormula_) {
		Expr left = elemFormula_.getLeft();
		Expr right = elemFormula_.getRight();
	
		if (elemFormula_.getOp().equals(CompOp.EQUALS) &&
		    (_transform(left,right) || _transform(right,left)))
		    return new EmptyFormula();
		    
		return elemFormula_;
    }

}
